from __future__ import division
from string import ascii_lowercase, maketrans
from collections import defaultdict
from operator import itemgetter
'''
Created on Jan 2, 2012

@author: Aaron Newton
'''

class RotationCipherSolver:
    '''
    Takes a string and solves for a Rotation (Caesar) Cipher
    For inspiration see http://code.activestate.com/recipes/442000-decoding-a-shift-or-rotation-or-caesar-cipher-or-c/
    '''

    def __init__(self, debugging=False, alphabet=ascii_lowercase, library_document="EnglishSample.txt", k_smoothing = 1):
        '''
        Kicks off the program - call newQuery to start __query_list
        '''
        'What __alphabet should we ue'
        self.__alphabet = alphabet
        self.__alphabet_list = tuple(alphabet) #','.join(string.ascii_lowercase)
        'Get the __query'
        self.__query = ""
        self.__query_list = []
        'Can be set to true to print out __debugging messages'
        self.__debugging = debugging #should we print out the various variables
        'Set the initial belief-state to unknown'
        self.__beliefState = []
        'Get the text from this document into memory'
        #self.library_document = open(library_document,'r')
        self.__library_text = self.__parse_library_book(library_document)
        #self.__library_text = self.library_document.readlines()
        'Analyse bigrams of __library_text'
        self.__library_bigrams = self.find_bigrams(self.__library_text) #gets the list of raw bigrams with duplicates
        self.__library_brigrams_scores = self.count_bigrams(self.__library_text) #generates a dictionary of unique bigrams with counts
        'scores for the states'
        self.__state_scores = {}
        'add a value for smoothing (laplace style)'
        self.__k_smoothing = k_smoothing #TODO - maybe we should use Katz Smoothing instead of +1 Smoothing - http://nlp.stanford.edu/~wcmac/papers/20050421-smoothing-tutorial.pdf
        #TODO - perhaps we could cache the above table
        if self.__debugging:
            print("Bigrams: " + str(self.__library_bigrams))
            print("Bigrams with counts: " + str(self.__library_brigrams_scores))
            
    def newQuery(self, query):
        '''
        Starts a new __query_list
        Takes single argument
        '''
        if self.__debugging:
            print("Query: " + query)
            #print("Query L_CASE:  " + __query.lower())
            print("Alphabet Constant: " + str(self.__alphabet_list))
        
        self.__query = query
        
        for character in query:
            self.__query_list.append(character)
            
        if self.__debugging:
            print("Query list: " + str(self.__query_list))
        'Generate the states'
        self.__brute_generate_states_and_store()
        'Attach scores to states'
        self.__state_scores = self.__score_states()
    
    def __parse_library_book(self,uri):
        """
        Get the URI/location of the text file, parse the book and return raw text for processing
        Basically serialize
        TODO - maximum efficiency string concatenation - see http://www.skymind.com/~ocrow/python_string/ method 6
        #http://stackoverflow.com/a/1019572/201648
        """
        content = ""
        book = open(uri,'r')
        line = book.readline()
        while line:
            line = line.rstrip() # because every line ends in a newline
            content += line + " "
            line = book.readline()
        return content
        
    
    def print_library_book(self):
        """
        Prints out the serialised content of the library book text
        """
        return self.__library_text

    def get_library_bigrams_counts(self):
        """
        Returns the rough bigram scores of the library bok text
        """
        return self.__library_brigrams_scores

    def __brute_generate_states_and_store(self):
        '''
        Takes the __alphabet_list, and for length of __alphabet_list increments the offset of the character by n
        The possible solution is then written to __beliefState
        '''
        
        'for the length of the __alphabet_list'
        '''
        for i in range(len(self.__alphabet_list)):
            print(i)
        '''
        
        alphaLength = len(self.__alphabet_list)
        loopIndices = range(alphaLength)
 
        "Generate all possible shifts"
        self.__beliefState = [self.shift_encode(self.__query, n) for n in range(len(self.__alphabet))] 
        
        if self.__debugging:
            print("Alphabet length: " + str(alphaLength))
            print("Indices: " + str(loopIndices))
            #print(self.__query)
            #print(self.shift_encode(self.__query, 5))
            #print(self.__beliefState)

    def get_belief_state(self):
        """
        Return out the belief state as a list
        """
        return self.__beliefState
    
    def print_belief_states(self):
        """
        Print belief states line by line
        """
        for state in self.__beliefState:
            print(state)
    
    def shift_encode(self, message, offset):
        """
        Takes a message and an offset, and performs a Caesar cipher shift for __alphabet_list on message
        Based on the Peter Norvig example
        See http://code.activestate.com/recipes/442000-decoding-a-shift-or-rotation-or-caesar-cipher-or-c/
        """
        
        'shift __alphabet_list by index n places'
        shifted_alphabet = self.__alphabet[offset:] + self.__alphabet[:offset] #take letters after offset, then stick the letters before offset on the end of this
        
        'take that __alphabet_list, compare it with original and develop a translation table'
        translator = maketrans(self.__alphabet + self.__alphabet.upper(), shifted_alphabet + shifted_alphabet.upper()) #http://www.tutorialspoint.com/python/string_maketrans.htm
        
        'now use the translator to translate the message'
        return message.translate(translator)
    
    def find_bigrams(self, text):
        """
        Counts bigrams (2 letter pairs, e.g. "ie", "ce", "ei")
        Duped from http://code.activestate.com/recipes/442000-decoding-a-shift-or-rotation-or-caesar-cipher-or-c/
        """
        return [text[i:i+2] for i in range(len(text) - 1)] #for each character, get that range of that character plus the next - return as a list
    
    def count_bigrams(self, text):
        """
        Gets the bigrams of the text and counts unique occurrences
        TODO - this is only an estimate: didn't count the exact number for 'In' (135) versus counted (127)
        We should probably check occurrences in the text instead
        BUT this is probably close enough for our purposes, and as heuristic should be admissible
        """
        "Get the bigrams"
        bigrams = self.find_bigrams(text)
        #http://stackoverflow.com/a/2392948/201648
        """
        c = Counter(bigrams)
        return c.items()
        """
        'Get as dictionary instead'
        #http://stackoverflow.com/a/2392986/201648
        d = defaultdict(int)
        for i in bigrams:
            d[i] += 1
        return d
    
    def __score_states(self):
        """
        Get the states and attch a score to each
        Private method/helper
        """
        scored_states = {}
        count_library_bigrams = len(self.__library_brigrams_scores)
        for state in self.__beliefState:
            bigrams = self.find_bigrams(state)
            score = 1.0
            count_state_bigrams = len(bigrams)
            for bigram in bigrams:
                #print(bigram)
                'How did this bigram score against the library book'
                bigram_library_score = self.__library_brigrams_scores[bigram]
                'I was having an issue where all states would have at least one "0" bigram answer, so I decided to apply smoothing'
                bigram_smoothed_score = (bigram_library_score + self.__k_smoothing) / (count_state_bigrams + count_library_bigrams)
                #self.__k_smoothing
                'If the score is larger than 0, multiply the score by this number'
                score = score * bigram_smoothed_score
                #print(bigram_smoothed_score, bigram_library_score, self.__k_smoothing, count_state_bigrams, count_library_bigrams, (bigram_library_score + self.__k_smoothing), (count_state_bigrams + count_library_bigrams))
                #print(self.__library_brigrams_scores[bigram])
            scored_states[state] = score
        return scored_states
    
    def get_state_scores(self):
        return self.__state_scores
    
    def get_state_by_max_score(self):
        """
        Once you have created a newQuery, you can call this to get the best decode by score
        """
        stats = self.__score_states()
        #http://stackoverflow.com/a/268285/201648
        return max(stats.iteritems(), key=itemgetter(1))[0]
        
        